# -*- coding: utf-8 -*-
import os
from locust import HttpUser, task, constant
from locust import LoadTestShape
from jsonpath import jsonpath


class QuickStartUser(HttpUser):
    host = "http://127.0.0.1:8000"
    wait_time = constant(1)

    @task
    def api(self):
        # print(time.strftime("%Y-%m-%d-%H-%M-%S"))
        with self.client.get(url="/locust_test", catch_response=True, name='测试') as response:
            # print(response.elapsed.total_seconds())
            # print(response.json())
            if response.status_code == 200:
                try:
                    assert jsonpath(response.json(), '$..code')[0] == 1
                except:
                    response.failure("Timeout")
            else:
                response.failure(response.status_code)


class StagesShape(LoadTestShape):
    """
    在不同的阶段 具有不同的用户数和 产生率的 图形形状
    参数解析:
        stages ：字典列表，每个字典都具有下列这些键值的阶段:
            duration -- 持续时间，  经过多少秒后，进入到下个阶段
            users -- 总用户数
            spawn_rate -- 产生率，即每秒启动/停止的用户数
            stop -- 可以在特定阶段停止测试的值
        stop_at_end -- 可以在所有阶段设置运行后停止

    """

    stages = [
        {"duration": 60, "users": 1, "spawn_rate": 1},
        {"duration": 120, "users": 3, "spawn_rate": 1},
        # {"duration": 180, "users": 5, "spawn_rate": 1},
        # {"duration": 240, "users": 7, "spawn_rate": 1},
        # {"duration": 300, "users": 9, "spawn_rate": 1},
        # {"duration": 360, "users": 11, "spawn_rate": 3},
        # {"duration": 420, "users": 13, "spawn_rate": 3},
        # {"duration": 480, "users": 15, "spawn_rate": 3},
        # {"duration": 540, "users": 17, "spawn_rate": 3},
        # {"duration": 600, "users": 19, "spawn_rate": 3},
    ]

    def tick(self):
        run_time = self.get_run_time()
        print(run_time)
        for stage in self.stages:
            if run_time < stage['duration']:
                tick_data = (stage['users'], stage['spawn_rate'])
                return tick_data

        return None


if __name__ == "__main__":
    os.system("locust -f 负载图_时间阶段负载策略.py --web-host=127.0.0.1")
